import torch
import numpy as np

# Constants are constructed with rdkit.Chem.GetPeriodicTable().GetAtomicWeight() functionality

NUM_TO_ELEMENT = {
    1: "H",
    2: "He",
    3: "Li",
    4: "Be",
    5: "B",
    6: "C",
    7: "N",
    8: "O",
    9: "F",
    10: "Ne",
    11: "Na",
    12: "Mg",
    13: "Al",
    14: "Si",
    15: "P",
    16: "S",
    17: "Cl",
    18: "Ar",
    19: "K",
    20: "Ca",
    21: "Sc",
    22: "Ti",
    23: "V",
    24: "Cr",
    25: "Mn",
    26: "Fe",
    27: "Co",
    28: "Ni",
    29: "Cu",
    30: "Zn",
    31: "Ga",
    32: "Ge",
    33: "As",
    34: "Se",
    35: "Br",
    36: "Kr",
    37: "Rb",
    38: "Sr",
    39: "Y",
    40: "Zr",
    41: "Nb",
    42: "Mo",
    43: "Tc",
    44: "Ru",
    45: "Rh",
    46: "Pd",
    47: "Ag",
    48: "Cd",
    49: "In",
    50: "Sn",
    51: "Sb",
    52: "Te",
    53: "I",
    54: "Xe",
    55: "Cs",
    56: "Ba",
    57: "La",
    58: "Ce",
    59: "Pr",
    60: "Nd",
    61: "Pm",
    62: "Sm",
    63: "Eu",
    64: "Gd",
    65: "Tb",
    66: "Dy",
    67: "Ho",
    68: "Er",
    69: "Tm",
    70: "Yb",
    71: "Lu",
    72: "Hf",
    73: "Ta",
    74: "W",
    75: "Re",
    76: "Os",
    77: "Ir",
    78: "Pt",
    79: "Au",
    80: "Hg",
    81: "Tl",
    82: "Pb",
    83: "Bi",
    84: "Po",
    85: "At",
    86: "Rn",
    87: "Fr",
    88: "Ra",
    89: "Ac",
    90: "Th",
    91: "Pa",
    92: "U",
}


ELEMENT_TO_NUM = {
    "H": 1,
    "He": 2,
    "Li": 3,
    "Be": 4,
    "B": 5,
    "C": 6,
    "N": 7,
    "O": 8,
    "F": 9,
    "Ne": 10,
    "Na": 11,
    "Mg": 12,
    "Al": 13,
    "Si": 14,
    "P": 15,
    "S": 16,
    "Cl": 17,
    "Ar": 18,
    "K": 19,
    "Ca": 20,
    "Sc": 21,
    "Ti": 22,
    "V": 23,
    "Cr": 24,
    "Mn": 25,
    "Fe": 26,
    "Co": 27,
    "Ni": 28,
    "Cu": 29,
    "Zn": 30,
    "Ga": 31,
    "Ge": 32,
    "As": 33,
    "Se": 34,
    "Br": 35,
    "Kr": 36,
    "Rb": 37,
    "Sr": 38,
    "Y": 39,
    "Zr": 40,
    "Nb": 41,
    "Mo": 42,
    "Tc": 43,
    "Ru": 44,
    "Rh": 45,
    "Pd": 46,
    "Ag": 47,
    "Cd": 48,
    "In": 49,
    "Sn": 50,
    "Sb": 51,
    "Te": 52,
    "I": 53,
    "Xe": 54,
    "Cs": 55,
    "Ba": 56,
    "La": 57,
    "Ce": 58,
    "Pr": 59,
    "Nd": 60,
    "Pm": 61,
    "Sm": 62,
    "Eu": 63,
    "Gd": 64,
    "Tb": 65,
    "Dy": 66,
    "Ho": 67,
    "Er": 68,
    "Tm": 69,
    "Yb": 70,
    "Lu": 71,
    "Hf": 72,
    "Ta": 73,
    "W": 74,
    "Re": 75,
    "Os": 76,
    "Ir": 77,
    "Pt": 78,
    "Au": 79,
    "Hg": 80,
    "Tl": 81,
    "Pb": 82,
    "Bi": 83,
    "Po": 84,
    "At": 85,
    "Rn": 86,
    "Fr": 87,
    "Ra": 88,
    "Ac": 89,
    "Th": 90,
    "Pa": 91,
    "U": 92,
}

ELEMENT_TO_MASS = {
    "H": 1.008,
    "He": 4.003,
    "Li": 6.941,
    "Be": 9.012,
    "B": 10.812,
    "C": 12.011,
    "N": 14.007,
    "O": 15.999,
    "F": 18.998,
    "Ne": 20.18,
    "Na": 22.99,
    "Mg": 24.305,
    "Al": 26.982,
    "Si": 28.086,
    "P": 30.974,
    "S": 32.067,
    "Cl": 35.453,
    "Ar": 39.948,
    "K": 39.098,
    "Ca": 40.078,
    "Sc": 44.956,
    "Ti": 47.867,
    "V": 50.942,
    "Cr": 51.996,
    "Mn": 54.938,
    "Fe": 55.845,
    "Co": 58.933,
    "Ni": 58.693,
    "Cu": 63.546,
    "Zn": 65.39,
    "Ga": 69.723,
    "Ge": 72.61,
    "As": 74.922,
    "Se": 78.96,
    "Br": 79.904,
    "Kr": 83.8,
    "Rb": 85.468,
    "Sr": 87.62,
    "Y": 88.906,
    "Zr": 91.224,
    "Nb": 92.906,
    "Mo": 95.94,
    "Tc": 98.0,
    "Ru": 101.07,
    "Rh": 102.906,
    "Pd": 106.42,
    "Ag": 107.868,
    "Cd": 112.412,
    "In": 114.818,
    "Sn": 118.711,
    "Sb": 121.76,
    "Te": 127.6,
    "I": 126.904,
    "Xe": 131.29,
    "Cs": 132.905,
    "Ba": 137.328,
    "La": 138.906,
    "Ce": 140.116,
    "Pr": 140.908,
    "Nd": 144.24,
    "Pm": 145.0,
    "Sm": 150.36,
    "Eu": 151.964,
    "Gd": 157.25,
    "Tb": 158.925,
    "Dy": 162.5,
    "Ho": 164.93,
    "Er": 167.26,
    "Tm": 168.934,
    "Yb": 173.04,
    "Lu": 174.967,
    "Hf": 178.49,
    "Ta": 180.948,
    "W": 183.84,
    "Re": 186.207,
    "Os": 190.23,
    "Ir": 192.217,
    "Pt": 195.078,
    "Au": 196.967,
    "Hg": 200.59,
    "Tl": 204.383,
    "Pb": 207.2,
    "Bi": 208.98,
    "Po": 209.0,
    "At": 210.0,
    "Rn": 222.0,
    "Fr": 223.0,
    "Ra": 226.0,
    "Ac": 227.0,
    "Th": 232.038,
    "Pa": 231.036,
    "U": 238.029,
}

MASS_TO_ELEMENT = {
    1.008: "H",
    4.003: "He",
    6.941: "Li",
    9.012: "Be",
    10.812: "B",
    12.011: "C",
    14.007: "N",
    15.999: "O",
    18.998: "F",
    20.18: "Ne",
    22.99: "Na",
    24.305: "Mg",
    26.982: "Al",
    28.086: "Si",
    30.974: "P",
    32.067: "S",
    35.453: "Cl",
    39.948: "Ar",
    39.098: "K",
    40.078: "Ca",
    44.956: "Sc",
    47.867: "Ti",
    50.942: "V",
    51.996: "Cr",
    54.938: "Mn",
    55.845: "Fe",
    58.933: "Co",
    58.693: "Ni",
    63.546: "Cu",
    65.39: "Zn",
    69.723: "Ga",
    72.61: "Ge",
    74.922: "As",
    78.96: "Se",
    79.904: "Br",
    83.8: "Kr",
    85.468: "Rb",
    87.62: "Sr",
    88.906: "Y",
    91.224: "Zr",
    92.906: "Nb",
    95.94: "Mo",
    98.0: "Tc",
    101.07: "Ru",
    102.906: "Rh",
    106.42: "Pd",
    107.868: "Ag",
    112.412: "Cd",
    114.818: "In",
    118.711: "Sn",
    121.76: "Sb",
    127.6: "Te",
    126.904: "I",
    131.29: "Xe",
    132.905: "Cs",
    137.328: "Ba",
    138.906: "La",
    140.116: "Ce",
    140.908: "Pr",
    144.24: "Nd",
    145.0: "Pm",
    150.36: "Sm",
    151.964: "Eu",
    157.25: "Gd",
    158.925: "Tb",
    162.5: "Dy",
    164.93: "Ho",
    167.26: "Er",
    168.934: "Tm",
    173.04: "Yb",
    174.967: "Lu",
    178.49: "Hf",
    180.948: "Ta",
    183.84: "W",
    186.207: "Re",
    190.23: "Os",
    192.217: "Ir",
    195.078: "Pt",
    196.967: "Au",
    200.59: "Hg",
    204.383: "Tl",
    207.2: "Pb",
    208.98: "Bi",
    209.0: "Po",
    210.0: "At",
    222.0: "Rn",
    223.0: "Fr",
    226.0: "Ra",
    227.0: "Ac",
    232.038: "Th",
    231.036: "Pa",
    238.029: "U",
}

EPSILON_0 = 8.854e-12  # C^2 * N^{-1} * m^{-2} (premittivity of vacuum) (F*m^{-1} - Farads per meters)
EPSILON_R = 1  # dielectric constant - set to 1 in vaccum
# COUL_CONST = 1/(4*torch.tensor(np.pi)*EPSILON_0)
COUL_CONST = 332.0636  # kcal*Angstrom/(mol*e^2) (https://www.ks.uiuc.edu/Research/namd//1.5/ug/node67.html)
HARTREE = 1.0
KJ_MOL = HARTREE / 2625.49963865
KCAL_MOL = HARTREE / 627.50947415
ANGSTROM = 1.0
BOHR = 0.52917721090380 * ANGSTROM
BOLTZMANN = 0.001988 * KCAL_MOL / HARTREE

SIGMA_ATOMS = {
    1: 2.5,
    2: 2.640,
    3: 2.103,
    4: 1.515,
    5: 3.58,
    6: 3.5,
    7: 3.2,
    8: 2.96,  # There are multiple values for oxygen
    9: 3.566,
    11: 2.87,
    12: 2.352,
    13: 4.05,
    14: 4.00,
    15: 3.74,
    16: 3.55,
    17: 4.504,
    19: 3.474,
    20: 3.12,
    35: 3.97,
    53: 5.448,
}

EPSILON_ATOMS = {
    1: 0.030,
    2: 0.18468,
    3: 0.026,
    4: 0.0043,
    5: 0.095,
    6: 0.066,
    7: 0.170,
    8: 0.21,
    9: 0.05,
    11: 0.0303,
    12: 0.034,
    13: 0.1,
    14: 0.1,
    15: 0.2,
    16: 0.25,
    17: 0.151,
    19: 0.0573,
    20: 0.073,
    35: 0.206,
    53: 0.2135,
}

CHARGES = {
    1: 0,
    2: 0,
    3: 0,
    4: 0,
    5: 0,
    6: 0,
    7: 0,
    8: 0,
    9: 0,
    11: 0,
    12: 0,
    13: 0,
    14: 0,
    15: 0,
    16: 0,
    17: 0,
    19: 0,
    20: 0,
    35: 0,
    53: 0,
}
